from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score, accuracy_score, precision_score, recall_score, f1_score
import xgboost as xgb
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve

class AnomalyDetector:
    def __init__(self):
        """初始化异常检测器"""
        self.lr_model = None
        self.xgb_model = None
        
    def train_models(self, X_train, y_train):
        """
        训练异常检测模型
        - 逻辑回归: 使用类权重平衡处理不平衡数据
        - XGBoost: 使用scale_pos_weight处理不平衡数据
        """
        print("开始训练异常检测模型...")
        
        # 计算正负样本比例
        pos_weight = len(y_train[y_train==0]) / len(y_train[y_train==1])
        
        # 训练逻辑回归
        print("训练逻辑回归模型...")
        self.lr_model = LogisticRegression(
            class_weight='balanced',
            max_iter=1000,
            random_state=42
        )
        self.lr_model.fit(X_train, y_train)
        
        # 训练XGBoost
        print("训练XGBoost模型...")
        self.xgb_model = xgb.XGBClassifier(
            scale_pos_weight=pos_weight,
            learning_rate=0.1,
            n_estimators=100,
            max_depth=5,
            random_state=42,
            eval_metric='auc'
        )
        self.xgb_model.fit(
            X_train, 
            y_train,
            eval_set=[(X_train, y_train)],
            verbose=False
        )
        
    def evaluate_models(self, X_test, y_test):
        """
        评估异常检测模型
        返回各种性能指标和ROC曲线
        """
        # 预测概率
        lr_proba = self.lr_model.predict_proba(X_test)[:, 1]
        xgb_proba = self.xgb_model.predict_proba(X_test)[:, 1]
        
        # 预测类别
        lr_pred = self.lr_model.predict(X_test)
        xgb_pred = self.xgb_model.predict(X_test)
        
        # 计算详细指标
        results = {
            'Logistic Regression': {
                'AUC-ROC': roc_auc_score(y_test, lr_proba),
                'Accuracy': accuracy_score(y_test, lr_pred),
                'Precision': precision_score(y_test, lr_pred),
                'Recall': recall_score(y_test, lr_pred),
                'F1-score': f1_score(y_test, lr_pred)
            },
            'XGBoost': {
                'AUC-ROC': roc_auc_score(y_test, xgb_proba),
                'Accuracy': accuracy_score(y_test, xgb_pred),
                'Precision': precision_score(y_test, xgb_pred),
                'Recall': recall_score(y_test, xgb_pred),
                'F1-score': f1_score(y_test, xgb_pred)
            }
        }
        
        # 绘制ROC曲线
        plt.figure(figsize=(10, 6))
        
        # 逻辑回归的ROC曲线
        fpr_lr, tpr_lr, _ = roc_curve(y_test, lr_proba)
        plt.plot(fpr_lr, tpr_lr, 
                label=f'Logistic Regression (AUC = {results["Logistic Regression"]["AUC-ROC"]:.3f})')
        
        # XGBoost的ROC曲线
        fpr_xgb, tpr_xgb, _ = roc_curve(y_test, xgb_proba)
        plt.plot(fpr_xgb, tpr_xgb, 
                label=f'XGBoost (AUC = {results["XGBoost"]["AUC-ROC"]:.3f})')
        
        plt.plot([0, 1], [0, 1], 'k--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('ROC Curves for Anomaly Detection Models')
        plt.legend(loc="lower right")
        plt.grid(True)
        plt.savefig('roc_curves.png')
        plt.close()
        
        return results

    def predict(self, X):
        """
        对新数据进行预测
        返回两个模型的预测结果
        """
        if self.lr_model is None or self.xgb_model is None:
            raise ValueError("模型尚未训练")
            
        results = {
            'Logistic Regression': {
                'predictions': self.lr_model.predict(X),
                'probabilities': self.lr_model.predict_proba(X)[:, 1]
            },
            'XGBoost': {
                'predictions': self.xgb_model.predict(X),
                'probabilities': self.xgb_model.predict_proba(X)[:, 1]
            }
        }
        
        return results 

    def evaluate_by_category(self, X_test, y_test, cell_ids, cat_ids):
        """按用户类别和RRU单元评估性能"""
        results = {}
        
        # 按RRU单元评估
        for cell_id in np.unique(cell_ids):
            cell_mask = cell_ids == cell_id
            if cell_mask.sum() > 0:
                results[f'Cell_{cell_id}'] = self.evaluate_models(
                    X_test[cell_mask], 
                    y_test[cell_mask]
                )
        
        # 按用户类别评估        
        for cat_id in np.unique(cat_ids):
            cat_mask = cat_ids == cat_id
            if cat_mask.sum() > 0:
                results[f'Category_{cat_id}'] = self.evaluate_models(
                    X_test[cat_mask],
                    y_test[cat_mask]
                )
            
        return results

    def plot_roc_curves(self, y_true, y_pred_proba, title):
        """绘制ROC曲线"""
        plt.figure(figsize=(10, 6))
        
        # 计算ROC曲线
        fpr, tpr, _ = roc_curve(y_true, y_pred_proba)
        roc_auc = roc_auc_score(y_true, y_pred_proba)
        
        # 绘制ROC曲线
        plt.plot(fpr, tpr, label=f'ROC curve (AUC = {roc_auc:.3f})')
        plt.plot([0, 1], [0, 1], 'k--')  # 随机猜测的基准线
        
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title(title)
        plt.legend(loc="lower right")
        plt.grid(True)
        plt.savefig(f'{title.lower().replace(" ", "_")}.png')
        plt.close() 